import matplotlib.pyplot as plt
import pandas as pd

dataset_path = './data/'
NUM_SAMPLES = 5000

if __name__ == "__main__":
    train_labels = pd.read_csv(dataset_path + '训练集label.csv')
    star_num = train_labels[train_labels['label'] == 'star'].shape[0]
    qso_num = train_labels[train_labels['label'] == 'qso'].shape[0]
    galaxy_num = train_labels[train_labels['label'] == 'galaxy'].shape[0]

    train_data = pd.read_csv(dataset_path + 'sets_v1_0.csv')
    plt.figure(figsize=(12, 9), dpi=100)
    plt.subplots_adjust(left=None, bottom=None, right=None, top=None, wspace=None, hspace=0.3)

    stars = train_labels[:NUM_SAMPLES][train_labels['label'] == 'star'].sample(n=3)
    qsos = train_labels[:NUM_SAMPLES][train_labels['label'] == 'qso'].sample(n=3)
    galaxys = train_labels[:NUM_SAMPLES][train_labels['label'] == 'galaxy'].sample(n=3)

    total = pd.concat([stars, qsos, galaxys])

    for i in range(9):
        plt.subplot(331 + i)
        fig_y = train_data.iloc[total.index[i], :train_data.shape[1]-1].values
        title = total['label'][total.index[i]]
        plt.title(title)
        plt.plot(fig_y, linestyle='-', color='grey')
    plt.show()

    labels = ['galaxy', 'star', 'qso']
    sizes = [galaxy_num, star_num, qso_num, ]
    colors = ['lightslategrey', 'sandybrown', 'yellowgreen']
    plt.pie(sizes, explode=(0.01, 0.02, 0.03), labels=labels, colors=colors, labeldistance=1.2,
            autopct='%3.2f%%', shadow=False, startangle=180, pctdistance=0.6)
    plt.axis('equal')
    plt.legend()
    plt.show()
